import math
from functools import lru_cache

import numpy as np
import torch
from torch import Tensor

from config import device


@lru_cache
def get_mean_vectors(sum: int, m: int):
    # 取m个数，数的和为sum
    if m == 1:
        return [[sum]]
    mean_vectors = []
    for i in range(1, sum - (m - 1) + 1):
        right_vec = get_mean_vectors(sum - i, m - 1)
        a = [i]
        for item in right_vec:
            mean_vectors.append(a + item)
    return mean_vectors


# 生成参考向量
@lru_cache
def get_ref_vectors(pop_size: int, m: int) -> Tensor:
    """
    生成均匀的参考向量
    :param pop_size:
    :param N: 每个维度上的间隔数。
    :param m: # 目标维度
    :return: 均匀的参考向量
    """
    N = get_n(pop_size, m)
    vectors = get_mean_vectors(N + m, m)
    vectors = (np.array(vectors) - 1) / N
    # print(len(vectors))   len = C m-1, N+m-1
    return torch.tensor(vectors, dtype=torch.double).round(decimals=5)


@lru_cache
def get_ref_vectors_grad(pop_size, m):
    """
    生成网格数据
    :param pop_size:
    :param m:
    :return:
    """
    gap = np.linspace(0, 1, int(np.ceil(pop_size ** (1 / m))))
    C = np.meshgrid(*([gap] * m))
    return torch.tensor(np.column_stack([c.reshape(-1) for c in C]), device=device)


def vector_radian(tensor0: Tensor, tensor1: Tensor, return_radian=True) -> Tensor:
    """
    计算两个二维张量的弧度值
    :param tensor0:
    :param tensor1:
    :param return_radian:为ture则返回夹角，false则返回cos值
    :return:
    """
    cos_sim = torch.cosine_similarity(tensor0.unsqueeze(dim=1), tensor1.unsqueeze(dim=0), dim=-1)
    return torch.acos(cos_sim) if return_radian else cos_sim


@lru_cache
def get_n(size: int, m: int) -> int:
    """
    用于计算get_ref_vectors方法的参数N
    :param size: 种群大小
    :param m: obj维度
    :return:
    """
    n = 1
    while math.comb(n + m - 1, m - 1) < size:
        n += 1
    if math.comb(n + m - 1, m - 1) > size:
        n = n - 1
    return n


def normalize_vector(v: Tensor, dim: int) -> Tensor:
    """
    向量单位化
    :param v:
    :param dim: 单位化的矩阵维度
    :return: 返回单位化的向量
    """
    norm = torch.norm(v, p=2, dim=dim, keepdim=True)
    unit_vector = v / norm
    return unit_vector


def unique(dec):
    """
    去重
    :param dec:
    :return:
    """
    temp = dec.clone()
    if dec.device.type == 'cuda':
        temp = dec.cpu()
    temp = temp.detach().numpy()
    _, index = np.unique(temp, axis=0, return_index=True)
    return torch.tensor(index, dtype=torch.int)
